// Copyright (c) Microsoft Corporation. All rights reserved. Licensed under the
// Microsoft Software License Terms for Microsoft Quantum Development Kit Libraries
// and Samples. See LICENSE in the project root for license information.

namespace Microsoft.Quantum.Research.Samples {
    open Microsoft.Quantum.Arrays;
    open Microsoft.Quantum.Intrinsic;
    open Microsoft.Quantum.Canon;
    open Microsoft.Quantum.Convert;
    open Microsoft.Quantum.Diagnostics;
    open Microsoft.Quantum.Math;
    open Microsoft.Quantum.Measurement;

    /// # Summary
    /// Applies the time evolution unitary generated by Hamiltonian
    /// H = X_1 + Z_1 Z_2, as simulated using a simple Trotter decomposition.
    operation EvolveByTrotter (time: Double, nStep: Int, qubits: Qubit[]) : Unit is Adj + Ctl {
        for (idxStep in 1..nStep) {
            Exp([PauliZ, PauliZ], -time / IntAsDouble(nStep), qubits);
            Exp([PauliX, PauliI], -time / IntAsDouble(nStep), qubits);
        }
    }

    /// # Summary
    /// Unitary that contains all information of a two-qubit Hamiltonian through ``qubitization.''
    /// # Description
    /// The unitary is $U = -i G S$ that encodes a Hamiltonian $H$ where $G = 2\ket{G}\bra{G} -1$ and $H = \bra{G} S \ket{G}$.
    operation ApplyBlockEncodedHamiltonian (qubits: Qubit[]) : Unit is Adj + Ctl {
        if (Length(qubits) < 3) {
            fail "Not enough qubits are allocated.";
        }

        let idxControl = 0;
        let idxPhysical1 = 1;
        let idxPhysical2 = 2;

        ExpFrac([PauliI], -1, 1, [qubits[idxControl]]); // global phase -i.
        ApplyWithCA(X, CX(_, qubits[idxPhysical1]), qubits[idxControl]); // control-X activated on |0> rather than |1>.
        ApplyToEachCA(CZ(qubits[idxControl], _), qubits[idxPhysical1..idxPhysical2]);
        X(qubits[idxControl]); // This is 2|G><G|-1, where |G> happens to be (|0> + |1>)/sqrt(2).
    }

    /// # Summary
    /// Converts a the unitary evolution applied by walkOperation to another
    /// operation, represented by a unitary operator whose eigenvalues are
    /// transformed by a function encoded in angles.
    ///
    /// # Input
    /// ## angles
    /// Angles measured in radians from x-axis of the Bloch sphere. These angles are computed by quantum signal preprocessing.
    /// There must be an odd number of angles.
    /// ## walkOperation
    /// A unitary operator, which can be anything but depends on the purpose.
    /// ## qubits
    /// A set of qubits on which walkOperation acts.
    /// # Remarks
    /// The eigenvalue transformation function that is realized in this operation by the input angles,
    /// always has even real part and odd imaginary part.
    /// The angles specify the directions of the projectors of primitive matrices. See [Haah, https://arxiv.org/abs/1806.10236].
    /// These are not exactly the angles that appear in [Low, Chuang, https://arxiv.org/abs/1610.06546], though they are linearly related.
    /// An auxillary qubit is allocated internally while running this operation.
    operation QuantumSignalProcessOnEquator(
        angles : Double[],
        walkOperation : (Qubit[] => Unit is Adj + Ctl),
        qubits : Qubit[]
    ) : Unit {
        let nQueries = Length(angles) - 1; // The number of times to call walkOperation.
        if (nQueries % 2 != 0) {
            fail "An odd number of calls to the quantum walk operator is forbidden.";
        }

        using (aux = Qubit()) {
            H(aux); // Preparing the auxillary qubit in the |+⟩ state.
            for (idx in 0..((nQueries / 2) - 1)) {
                within {
                    Rz(-angles[2 * idx], aux);
                    H(aux);
                    X(aux);
                } apply {
                    (Controlled walkOperation)([aux], qubits);
                }

                within {
                    Rz(-angles[2 * idx + 1], aux);
                    H(aux);
                } apply {
                    Controlled Adjoint walkOperation([aux], qubits);
                }

            }
            Rz(angles[nQueries], aux); // The last unitary correction.
            if (MResetX(aux) == One) {
                fail "QSP's postselection failed.";
            }
        }
    }

    /// # Summary
    /// We use quantum signal processing to implement the time evolution unitary by a Hamiltonian on two qubits.
    /// An initial state is evolved forward by quantum signal processing, and then backward by Trotter steps.
    /// The angles for quantum signal processing depends only on tau and accuracy parameter.
    /// If the accuracy parameter is sufficiently small, say, 1.0e-20, then
    /// the initial state and the final state are the same up to error that is dominated by the Trotter step size,
    /// which is in this case about 1.0e-3.
    /// # Remark
    /// The parameter tau is the evolution time multiplied by the 1-norm of the coefficient of the Hamiltonian, which is 2 in this sample.
    /// The angles is precomputed by a quantum signal processing module,
    /// which should be called beforehand by a C# driver.
    operation SampleHamiltonianEvolutionByQSP(tau : Double, angles : Double[])
    : Unit {
        using (qubits = Qubit[3]) {
            H(Head(qubits)); // Preparing an auxillary qubit in the state $\ket{G}$.
            DumpMachine(());
            QuantumSignalProcessOnEquator(angles, ApplyBlockEncodedHamiltonian, qubits);
            // Now we roll back by Trotter evolution.
            EvolveByTrotter(-tau / 2.0, 1000, qubits[1..2]);
            // The factor of 2 that divides tau is the 1-norm of the Hamiltonian's coefficients.
            DumpMachine(());
            ResetAll(qubits);
        }
    }
}
